import os
import os.path as osp
import json
import argparse
from typing import List, Dict

import pandas as pd


def extract_category(caption: str) -> str:
    """Extract disease category from caption.
    Rules:
    - If contains 'healthy' (case-insensitive), return 'healthy'.
    - If contains 'diseased by <name>', return <name> trimmed, and cut suffix ' with symptoms ...'.
    - Else return 'unknown'.
    """
    if not isinstance(caption, str):
        return 'unknown'
    lower = caption.lower()
    if 'healthy' in lower:
        return 'healthy'
    key = 'diseased by'
    idx = lower.find(key)
    if idx >= 0:
        start = idx + len(key)
        after = caption[start:].strip()
        cut_mark = ' with symptoms'
        cut_idx = after.lower().find(cut_mark)
        if cut_idx >= 0:
            after = after[:cut_idx]
        return ' '.join(after.split())
    return 'unknown'


def extract_crop(caption: str) -> str:
    """Extract crop name from caption of the form:
    'a image of <Crop> healthy leaves ...' or 'a image of <Crop> leaves diseased by ...'
    Returns lowercase-trimmed crop name with original case preserved best-effort.
    """
    if not isinstance(caption, str):
        return ''
    text = caption.strip()
    lower = text.lower()
    key = 'a image of '
    idx = lower.find(key)
    if idx < 0:
        return ''
    start = idx + len(key)
    # Try " healthy leaves" first to avoid capturing healthy as crop
    end_markers = [' healthy leaves', ' leaves']
    end = -1
    for mark in end_markers:
        j = lower.find(mark, start)
        if j >= 0:
            end = j
            break
    if end < 0:
        # fallback: up to first space then word 'leaves' if present later
        end = lower.find(' leaves', start)
        if end < 0:
            return ''
    crop = text[start:end].strip()
    # normalize spaces
    crop = ' '.join(crop.split())
    return crop


def build_label_mapping(categories: List[str]) -> Dict[str, int]:
    uniq = sorted(set(categories))
    return {c: i for i, c in enumerate(uniq)}


def main():
    parser = argparse.ArgumentParser(description='Convert metadata.json to zero-shot CSV')
    parser.add_argument('--metadata', required=True, help='Path to metadata.json')
    parser.add_argument('--out_csv', required=True, help='Output CSV path for zero-shot (image[,label])')
    parser.add_argument('--image_root', required=True, help='Image root directory used by zeroshot_eval')
    parser.add_argument('--include_labels', action='store_true', help='Include integer labels derived from captions')
    parser.add_argument('--include_text', action='store_true', help='Include per-sample text column from caption')
    parser.add_argument('--split_healthy_by_crop', action='store_true', help='Split healthy into <Crop> healthy classes')
    parser.add_argument('--split_diseases_by_crop', action='store_true', help='Split disease classes into <Crop> <Disease>')
    parser.add_argument('--classes_out', default=None, help='Optional classes.txt output when include_labels is set')
    parser.add_argument('--label_map_out', default=None, help='Optional label2id.json output when include_labels is set')
    parser.add_argument('--check_paths', action='store_true', help='Sanity check image existence under image_root')
    args = parser.parse_args()

    if not osp.exists(args.metadata):
        raise FileNotFoundError(f"metadata.json not found: {args.metadata}")
    if not osp.isdir(args.image_root):
        raise FileNotFoundError(f"image_root not found: {args.image_root}")

    with open(args.metadata, 'r', encoding='utf-8') as f:
        data = json.load(f)

    rows = []
    categories: List[str] = []
    for d in data:
        # Prefer image_filename as relative path under image_root
        img_rel = d.get('image_filename')
        if not img_rel:
            file_name = str(d.get('file_name', ''))
            if file_name:
                img_rel = osp.basename(file_name)
        if not img_rel:
            img_path = str(d.get('image_path', ''))
            if img_path:
                img_rel = osp.basename(img_path)
        if not img_rel:
            # skip entries without resolvable image name
            continue

        row = {'image': img_rel}
        caption = d.get('caption', '') or ''
        if args.include_labels:
            base_cat = extract_category(caption)
            crop = extract_crop(caption) if (args.split_healthy_by_crop or args.split_diseases_by_crop) else ''
            cat = base_cat
            # apply splitting rules
            if args.split_healthy_by_crop and base_cat.lower() == 'healthy' and crop:
                cat = f"{crop} healthy"
            elif args.split_diseases_by_crop and base_cat.lower() != 'unknown' and base_cat.lower() != 'healthy' and crop:
                cat = f"{crop} {base_cat}"
            categories.append(cat)
            row['disease_category'] = cat
        if args.include_text:
            # use caption as text
            row['text'] = caption
        rows.append(row)

    df = pd.DataFrame(rows)
    if df.empty:
        raise ValueError('No valid entries found in metadata.json to build zero-shot CSV')

    if args.include_labels:
        label2id = build_label_mapping(categories)
        df['label'] = df['disease_category'].map(label2id)

        # Save mapping artifacts
        if args.label_map_out:
            id2label = {v: k for k, v in label2id.items()}
            with open(args.label_map_out, 'w', encoding='utf-8') as f:
                json.dump({'label2id': label2id, 'id2label': id2label}, f, ensure_ascii=False, indent=2)
        if args.classes_out:
            inv = sorted(label2id.items(), key=lambda kv: kv[1])
            with open(args.classes_out, 'w', encoding='utf-8') as f:
                for c, _id in inv:
                    f.write(f"{c}\n")

    # Optional path existence check
    if args.check_paths:
        exists = df['image'].apply(lambda fn: osp.exists(osp.join(args.image_root, str(fn))))
        missing = (~exists).sum()
        if missing:
            print(f"[warn] {missing} images not found under {args.image_root}")

    # Write CSV
    os.makedirs(osp.dirname(args.out_csv), exist_ok=True)
    df.to_csv(args.out_csv, index=False)
    print(f"Wrote zero-shot CSV to {args.out_csv} (rows={len(df)})")

    if args.include_labels:
        dist = df['disease_category'].value_counts().sort_values(ascending=False)
        dist_path = osp.join(osp.dirname(args.out_csv), 'class_distribution.csv')
        dist.to_csv(dist_path)
        print(f"Saved class_distribution.csv to {dist_path}")


if __name__ == '__main__':
    main()